module tds4py
  use constants
  use parameters
  use basic
  use input_output
  use grt_coefficients
  use matrix_operations
  use wavefield_calculation
  use, intrinsic :: iso_c_binding
  implicit none

contains

  subroutine init_tds_parameters()
    call init_all_parameters()
  end subroutine init_tds_parameters

  subroutine debug(n_layers, z, vp0, vs0, rho, qs, qp)
    integer, intent(in) :: n_layers
    real(kind=8), intent(in) :: z(0:4), vp0(4), vs0(4), &
                               rho(4), qs(4), qp(4)
    integer :: i
    print *,n_layers
    print *,z
    print *,vp0
    print *,vs0

  end subroutine debug


  subroutine set_media_parameters(n_layers, z, vp0, vs0, rho, qs, qp)
    integer, intent(in) :: n_layers
    real(kind=8), intent(in) :: z(0:n_layers), vp0(n_layers), vs0(n_layers), &
                               rho(n_layers), qs(n_layers), qp(n_layers)
    integer :: i

    nly = n_layers
    
    if (allocated(media%z)) deallocate(media%z)
    if (allocated(media%vp0)) deallocate(media%vp0)
    if (allocated(media%vs0)) deallocate(media%vs0)
    if (allocated(media%rho)) deallocate(media%rho)
    if (allocated(media%qs)) deallocate(media%qs)
    if (allocated(media%qp)) deallocate(media%qp)
    if (allocated(media%mu)) deallocate(media%mu)
    if (allocated(media%vs)) deallocate(media%vs)
    if (allocated(media%vp)) deallocate(media%vp)
    
    allocate(media%z(0:nly))
    allocate(media%vp0(nly))
    allocate(media%vs0(nly))
    allocate(media%rho(nly))
    allocate(media%qs(nly))
    allocate(media%qp(nly))
    allocate(media%mu(nly))
    allocate(media%vs(nly))
    allocate(media%vp(nly))
    
    do i = 0, nly
      media%z(i) = z(i)
    end do
    
    do i = 1, nly
      media%vp0(i) = vp0(i)
      media%vs0(i) = vs0(i)
      media%rho(i) = rho(i)
      media%qs(i) = qs(i)
      media%qp(i) = qp(i)
    end do
  end subroutine set_media_parameters

  subroutine set_source_parameters(zs, z0, x0, y0, xs, ys, fc, tou, M0, dip, rake, str, &
                                  mino, maxo, no, minv, maxv, nv)
    real(kind=8), intent(in) :: zs, z0, x0, y0, xs, ys, fc, tou, M0, dip, rake, str, &
                               mino, maxo, minv, maxv
    integer, intent(in) :: no, nv
    
    source%zs = zs
    source%z0 = z0
    
    source%r0 = sqrt((x0-xs)**2 + (y0-ys)**2)
    if (y0 >= ys) then
      source%fai0 = acos((x0-xs)/source%r0)
    else
      source%fai0 = pi2 - acos((x0-xs)/source%r0)
    end if
    
    source%fc = fc
    source%tou = tou
    source%M0 = M0
    source%dip = dip * pi / 180.0d0  ! 角度转弧度
    source%rake = rake * pi / 180.0d0  ! 角度转弧度
    source%str = str * pi / 180.0d0  ! 角度转弧度
    source%mino = mino
    source%maxo = maxo
    source%no = no
    source%minv = minv
    source%maxv = maxv
    source%nv = nv
    
    ls = layernumb(source%zs, nly, media%z)
    lo = layernumb(source%z0, nly, media%z)
  end subroutine set_source_parameters

  ! 设置计算参数
  subroutine set_computation_parameters(m, Twin, oi)
    integer, intent(in) :: m
    real(kind=8), intent(in) :: Twin, oi
    
    comp%m = m
    comp%Twin = Twin
    comp%oi = oi
  end subroutine set_computation_parameters

  ! 设置其他参数
  subroutine set_other_parameters(SourceType_val, Case_val, Type_STF_val, &
                                 Outname_val, OutFormat_val, WinSwitch_val, &
                                 WinType_val, f1_val, f2_val, f3_val, f4_val)
    character(len=*), intent(in) :: SourceType_val, Case_val, Type_STF_val, &
                                   Outname_val, OutFormat_val, WinSwitch_val, &
                                   WinType_val
    real(kind=8), intent(in) :: f1_val, f2_val, f3_val, f4_val
    
    SourceType = SourceType_val
    Case = Case_val
    Type_STF = Type_STF_val
    Outname = Outname_val
    OutFormat = OutFormat_val
    WinSwitch = WinSwitch_val
    WinType = WinType_val
    f1 = f1_val
    f2 = f2_val
    f3 = f3_val
    f4 = f4_val
  end subroutine set_other_parameters

  ! 计算频散图（使用频率和速度范围）
  subroutine calculate_dispersion_spectrum(mino,maxo,minv, maxv, no,nv, result)
    integer, intent(in) :: no, nv
    real(kind=8), intent(in) :: mino, maxo, minv, maxv
    complex(kind=8), intent(out) :: result(no, nv)
    
    real(kind=8), allocatable :: Vph(:)        ! 相速度数组
    complex(kind=8), allocatable :: omega(:)   ! 角频率数组
    complex(kind=8) :: integ(10)               ! 积分结果
    real(kind=8) :: kn 
    complex(kind=8) :: o                  ! 波数和角频率
    complex(kind=8) :: cs, cp                 ! 复速度
    integer :: i, j, ii

    ! 分配数组内存
    allocate(omega(no))
    allocate(Vph(nv))

    call green_basic()

    ! 构建频率和相速度数组
    do i = 1, no
      omega(i) = (mino + (i-1)*(maxo-mino)/(no-1.0d0)) * pi2
    end do

    do i = 1, nv
      Vph(i) = minv + (i-1)*(maxv-minv)/(nv-1.0d0)
    end do

    ! 主计算循环
    do i = 1, no
      ! 计算复速度
      do ii = 1, nly
        cs = (1.0d0 + log(omega(i)/(pi2*maxo))/(pi*media%qs(ii)) + aj/(2.0d0*media%qs(ii)))
        cp = (1.0d0 + log(omega(i)/(pi2*maxo))/(pi*media%qp(ii)) + aj/(2.0d0*media%qp(ii)))
        media%vs(ii) = media%vs0(ii) * cs
        media%vp(ii) = media%vp0(ii) * cp
      end do
      
      ! 对每个相速度计算响应
      do j = 1, nv
        kn = omega(i) / Vph(j)
        o = omega(i) - aj * comp%oi

        ! 计算位移响应
        call funval(kn, o, integ)
        ! 保存结果（使用瑞利波垂向位移分量）
        result(i, j) = integ(10)
      end do
    end do

    ! 释放内存
    deallocate(omega, Vph)
  end subroutine calculate_dispersion_spectrum

  ! 计算频散图（使用频率和速度数组）
  subroutine calculate_dispersion_spectrum_arrays(freq_array, vel_array, np, result)
    integer, intent(in) :: np
    real(kind=8), intent(in) :: freq_array(np), vel_array(np)
    complex(kind=8), intent(out) :: result(np)
    
    complex(kind=8) :: omega_val
    real(kind=8) :: kn
    real(kind=8) :: maxo
    complex(kind=8) :: integ(10)  ! 积分结果
    complex(kind=8) :: o          ! 波数和角频率
    complex(kind=8) :: cs, cp     ! 复速度
    integer :: i, j, ii

    ! 计算omega_val的最大值
    maxo = maxval(freq_array)

    call green_basic()
    

    ! 主计算循环
    do i = 1, np
      omega_val = freq_array(i) * pi2
      ! print *, 'f:', freq_array(i), 'v:',vel_array(i)
      ! 计算复速度
      do ii = 1, nly
        cs = (1.0d0 + log(omega_val/(pi2*maxo))/(pi*media%qs(ii)) + aj/(2.0d0*media%qs(ii)))
        cp = (1.0d0 + log(omega_val/(pi2*maxo))/(pi*media%qp(ii)) + aj/(2.0d0*media%qp(ii)))
        media%vs(ii) = media%vs0(ii) * cs
        media%vp(ii) = media%vp0(ii) * cp
      end do
      
      ! 对每个相速度计算响应
      kn = omega_val / vel_array(i)
      o = omega_val - aj * comp%oi

      ! 计算位移响应
      call funval(kn, o, integ)
      ! 保存结果（使用瑞利波垂向位移分量）
      result(i) = integ(10)
      
    end do
  end subroutine calculate_dispersion_spectrum_arrays

  ! 保存参数到grt.dat文件，方便与原始main.f90结果对比验证
  subroutine save_grt_dat()
    
    integer :: file_unit = 10
    integer :: i
    
    open(file_unit, file=trim(Outname)//'.grt1', form='formatted', status='replace')
    
    ! 写入接收器参数
    write(file_unit, '(A)') 'FOR : receiver information: rarray case:'
    write(file_unit, '(A)') ''
    write(file_unit, '(F10.2)') 0.0d0
    write(file_unit, '(F10.2)') 0.0d0
    write(file_unit, '(F10.2)') source%z0
    
    ! 写入源参数
    write(file_unit, '(A)') ''
    write(file_unit, '(A)') 'FOR: input dat for double-couple source'
    write(file_unit, '(A)') ''
    write(file_unit, '(E25.2)') source%M0
    write(file_unit, '(F6.4,2X,F6.2)') source%mino, source%maxo
    write(file_unit, '(I6)') source%no
    write(file_unit, '(F6.2,2X,F6.2)') source%minv, source%maxv
    write(file_unit, '(I6)') source%nv
    write(file_unit, '(3(F6.2,1X))') source%rake*180.0d0/pi, source%dip*180.0d0/pi, source%str*180.0d0/pi
    write(file_unit, '(F6.2)') 0.0d0
    write(file_unit, '(F6.2)') 0.0d0
    write(file_unit, '(F6.2)') source%zs
    
    ! 写入计算参数
    write(file_unit, '(A)') ''
    write(file_unit, '(A)') 'FOR: basic control parameters'
    write(file_unit, '(A)') ''
    write(file_unit, '(A)') '(1). Basic control parameters'
    write(file_unit, '(I6)') comp%m
    write(file_unit, '(F12.6)') comp%Twin
    
    ! 写入源类型
    write(file_unit, '(A)') ''
    write(file_unit, '(A)') '(2).Source Typer'
    write(file_unit, '(A)') trim(SourceType)
    
    ! 写入源谱参数
    write(file_unit, '(A)') ''
    write(file_unit, '(A)') '(3)Source spectral parameters'
    write(file_unit, '(A)') trim(Type_STF)
    write(file_unit, '(A)') ''
    write(file_unit, '(F6.2)') source%tou
    write(file_unit, '(F6.2)') 0.0d0  ! shift
    write(file_unit, '(F6.2)') source%fc
    
    ! 写入窗口参数
    write(file_unit, '(A)') ''
    write(file_unit, '(A)') '(4).Window'
    write(file_unit, '(A)') trim(WinSwitch)
    write(file_unit, '(A)') trim(WinType)
    write(file_unit, '(4(F6.2,1X))') f1, f2, f3, f4
    
    ! 写入输出格式
    write(file_unit, '(A)') ''
    write(file_unit, '(A)') '(5).Out Format'
    write(file_unit, '(A)') trim(Outname)
    
    ! 写入媒质参数
    write(file_unit, '(A)') ''
    write(file_unit, '(A)') 'FOR: media'
    write(file_unit, '(A)') 'Number of total layers'
    write(file_unit, '(I6)') nly
    write(file_unit, '(A)') 'Layer ...'
    
    ! 写入各层参数
    do i = 1, nly
      write(file_unit, '(I4,2X,F8.4,2X,F8.4,2X,F8.4,2X,F8.4,2X,I6,2X,I6)') &
        i, media%z(i-1), media%rho(i), media%vs0(i), media%vp0(i), int(media%qs(i)), int(media%qp(i))
    end do
    
    close(file_unit)
  end subroutine save_grt_dat

end module tds4py